# asr_denoise.py
"""
ASR 噪声去除模块（Denoising ASR Noise）
功能：从 JSON 文件加载噪声-标准映射，还原 ASR 错误表达
"""

import json
import os
from typing import Dict, List
import re



asr_noise_map_path = "../data/asr_noise_map.json"
class ASRDenormalizer:
    """
    ASR 噪声还原器
    从外部 JSON 文件加载噪声-标准映射
    """

    def __init__(self, noise_map_file: str = None, custom_noise_map: Dict[str, str] = None):
        """
        初始化
        :param noise_map_file: 噪声映射 JSON 文件路径
                             默认为同目录下 config/asr_noise_map.json
        :param custom_noise_map: 用户自定义映射（可选，用于补充或覆盖）
        """
        # 设置默认文件路径
        if noise_map_file is None:
            noise_map_file = asr_noise_map_path

        # 加载 JSON 映射表
        self.noise_map = self._load_noise_map_from_json(noise_map_file)

        # 合并用户传入的自定义映射（可覆盖）
        if custom_noise_map:
            self.noise_map.update(custom_noise_map)

        # 编译正则表达式用于高效匹配
        self._compiled_pattern = self._compile_pattern()

    def _load_noise_map_from_json(self, file_path: str) -> Dict[str, str]:
        """从 JSON 文件加载噪声映射"""
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"噪声映射文件未找到: {file_path}")

        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                noise_map = json.load(f)

            # 验证数据类型
            if not isinstance(noise_map, dict):
                raise ValueError(f"噪声映射文件必须是一个 JSON 对象（dict），实际类型: {type(noise_map)}")

            # 过滤掉值为空的项
            noise_map = {k: v for k, v in noise_map.items() if k.strip() and v}

            print(f"成功加载 {len(noise_map)} 条噪声映射规则: {file_path}")
            return noise_map

        except json.JSONDecodeError as e:
            raise ValueError(f"JSON 文件格式错误: {file_path}, 错误: {e}")

    def _compile_pattern(self) -> 're.Pattern':
        """编译正则表达式，按长度降序匹配（避免部分匹配）"""
        # 按键长度降序排列，确保长字符串优先匹配
        sorted_noises = sorted(self.noise_map.keys(), key=len, reverse=True)
        pattern = '|'.join(re.escape(noise) for noise in sorted_noises)
        return re.compile(pattern)

    def denormalize(self, text: str) -> str:
        """
        去除 ASR 噪声，还原为标准文本
        :param text: 原始含噪声文本
        :return: 去噪后的标准文本
        """
        if not isinstance(text, str) or not text.strip():
            return text

        def replace_match(match):
            return self.noise_map.get(match.group(0), match.group(0))

        return self._compiled_pattern.sub(replace_match, text)

    def denormalize_batch(self, texts: List[str]) -> List[str]:
        """
        批量去噪
        :param texts: 文本列表
        :return: 去噪后的文本列表
        """
        return [self.denormalize(text) for text in texts]

    def add_noise_mapping(self, noise: str, standard: str):
        """
        动态添加新的噪声-标准映射
        :param noise: 噪声表达
        :param standard: 标准表达
        """
        self.noise_map[noise] = standard
        # 重新编译正则表达式
        self._compiled_pattern = self._compile_pattern()

    def get_all_mappings(self) -> Dict[str, str]:
        """返回当前所有映射（用于调试或保存）"""
        return self.noise_map.copy()


# ======================================
# 使用示例
# ======================================

if __name__ == "__main__":
    # 使用默认路径
    denormalizer = ASRDenormalizer()

    # 测试
    noisy_texts = [
        "帮我拨放一首周节伦的歌",
        "把空条温度调高",
        "热水瓶坏了",
        "脑中响了",
        "七点叫我起床",
        "风束调大一点",
        "音浪静音",
        "查一下茶询航班"
    ]

    print("ASR 去噪测试：")
    for noisy in noisy_texts:
        cleaned = denormalizer.denormalize(noisy)
        print(f"  '{noisy}' → '{cleaned}'")